package org.ws.task.repository;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang.StringUtils;
import org.apache.curator.shaded.com.google.common.collect.Sets;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.ws.task.item.TaskInstance;
import org.ws.task.model.NodeMetadata;
import org.ws.task.model.TaskConfig;

import java.util.Collection;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicReference;

public class NodeMetadataRepository {

    private static Logger log = LoggerFactory.getLogger(NodeMetadataRepository.class.getName());

    private static final NodeMetadataRepository instance = new NodeMetadataRepository();
    private final AtomicReference<NodeMetadata> leaderNode;
    private final AtomicReference<NodeMetadata> theWorkerNode;
    private final Map<String, Set<String>> activeWorkerNodeSetMap;
    private final Map<String,NodeMetadata> workerNodeMap;

    private NodeMetadataRepository(){
        this.leaderNode = new AtomicReference<>();
        this.theWorkerNode = new AtomicReference<>();
        this.activeWorkerNodeSetMap = Maps.newHashMap();
        this.workerNodeMap = Maps.newHashMap();
    }

    public static NodeMetadataRepository getInstance(){
        return instance;
    }

    public NodeMetadata getLeader(){
        return this.leaderNode.get();
    }

    public void changeLeader(NodeMetadata leaderNode){
        this.leaderNode.set(leaderNode);
    }

    public void setTheWorkerNode(NodeMetadata workerNode){
        this.theWorkerNode.set(workerNode);
    }

    public NodeMetadata getTheWorkerNode(){
        return this.theWorkerNode.get();
    }

    /**
     * 根据ID更新
     * @param workerNode
     */
    public synchronized void updateWorker(NodeMetadata workerNode){
        // 更新节点
        this.workerNodeMap.put(workerNode.getId(),workerNode);
        // 本节点
        if(this.theWorkerNode.get() != null
                && StringUtils.equals(workerNode.getId(),this.theWorkerNode.get().getId())){
            this.theWorkerNode.set(workerNode);
        }
        // 映射关系
        mappingTaskToWorkerNode(workerNode);
        // 无效状态直接删除
        if(!workerNode.isActive()){
            Collection<Set<String>> values = this.activeWorkerNodeSetMap.values();
            if(CollectionUtils.isNotEmpty(values)) {
                for (Set<String> value : values) {
                    if(value.contains(workerNode.getId())){
                        value.remove(workerNode.getId());
                    }
                }
            }
        }
    }

    public NodeMetadata getNode(String id){
        return this.workerNodeMap.get(id);
    }

    private void mappingTaskToWorkerNode(NodeMetadata workerNode){
        Set<String> taskNames = TaskRepository.getInstance().taskNames();
        for (String taskName : taskNames) {
            List<NodeMetadata> activeWorkerNodes = this.getActiveWorkerNodes(taskName);
            int activeNodeSize = activeWorkerNodes.size();
            TaskInstance taskInstance = TaskRepository.getInstance().getTaskInstance(taskName);
            TaskConfig taskConfig = taskInstance.getTaskConfig();
            Integer workerNum = taskConfig.getWorkerNum();
            Boolean leaderProcess = taskConfig.getLeaderProcess();
            NodeMetadata leader = NodeMetadataRepository.getInstance().getLeader();
            // Leader不处理任务 && Leader
            if(!leaderProcess && leader != null){
                // 当前为leader && 已存在worker -> 不分配
                if(StringUtils.equals(leader.getId(),workerNode.getId()) && activeNodeSize > 0){
                    log.warn("not mapping taskName node:{} -> leaderProcess:{},isLeader:{},activeNodeSize:{}",taskName,leaderProcess,workerNode.isLeader(),activeNodeSize);
                    continue;
                }
                // 当前位worker -> 移除leader
                if(!StringUtils.equals(leader.getId(),workerNode.getId())){
                    Set<String> ids = this.activeWorkerNodeSetMap.get(taskName);
                    if(CollectionUtils.isNotEmpty(ids)){
                        ids.remove(leader.getId());
                    }
                }
                // 当前无效 && 没有有效的node，加入Leader
                else if(!workerNode.isActive() && activeNodeSize<=0){
                    Collection<NodeMetadata> workers = NodeMetadataRepository.getInstance().getWorkers();
                    for (NodeMetadata worker : workers) {
                        if(StringUtils.equals(leader.getId(),worker.getId())){
                            addTaskWorkerNode(taskName,worker);
                        }
                    }
                }
            }
            if(workerNum<=0
                    || activeNodeSize < workerNum){
                addTaskWorkerNode(taskName,workerNode);
                log.debug("mapping taskName node:{} -> workerNum:{},activeNodeSize:{}",taskName,workerNum,activeNodeSize);
            }
            else{
                log.warn("not mapping taskName node:{} -> workerNum:{},activeNodeSize:{}",taskName,workerNum,activeNodeSize);
            }
        }
    }

    /***
     * 添加任务与处理节点关系
     * @param taskName
     * @param workerNode
     */
    private void addTaskWorkerNode(String taskName,NodeMetadata workerNode){
        if (!this.activeWorkerNodeSetMap.containsKey(taskName)) {
            this.activeWorkerNodeSetMap.putIfAbsent(taskName, Sets.newConcurrentHashSet());
        }
        if(workerNode.isActive()) {
            Set<String> ids = this.activeWorkerNodeSetMap.get(taskName);
            ids.add(workerNode.getId());
        }
    }

    /**
     * 根据id移除以及对应关系
     * @param workerNode
     */
    public synchronized void removeWorker(NodeMetadata workerNode){
        String id = workerNode.getId();
        this.workerNodeMap.remove(id);
        Collection<Set<String>> values = this.activeWorkerNodeSetMap.values();
        if(CollectionUtils.isNotEmpty(values)){
            for (Set<String> value : values) {
                if(value.contains(id)){
                    value.remove(id);
                }
            }
        }
    }

    public Collection<NodeMetadata> getWorkers(){
        return this.workerNodeMap.values();
    }


    public synchronized Set<String> getTaskWorkerNodeIds(String taskName){
        if (!this.activeWorkerNodeSetMap.containsKey(taskName)) {
            return Sets.newHashSet();
        }
        return this.activeWorkerNodeSetMap.get(taskName);
    }

    /**
     * 获取任务处理节点
     * @param taskName
     * @return
     */
    public List<NodeMetadata> getActiveWorkerNodes(String taskName){
        List<NodeMetadata> nodeMetadataList = Lists.newArrayList();
        Set<String> ids = this.getTaskWorkerNodeIds(taskName);
        for (String id : ids) {
            NodeMetadata nodeMetadata = this.workerNodeMap.get(id);
            nodeMetadataList.add(nodeMetadata);
        }
        return nodeMetadataList;
    }

    /**
     * 节点能否处理任务
     * @param taskName
     * @param nodeId
     * @return
     */
    public synchronized boolean isMatch(String taskName,String nodeId){
        if (!this.activeWorkerNodeSetMap.containsKey(taskName)) {
            return false;
        }
        Set<String> ids = this.activeWorkerNodeSetMap.get(taskName);
        return ids.contains(nodeId);
    }

}
